/**
 *   Copyright 2009 Class Action P/L, Australia
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 */

package com.classactionpl.jaas.activemq;

import java.security.cert.X509Certificate;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Hashtable;
import java.util.Map;
import java.util.Set;

import javax.naming.Context;
import javax.naming.NamingEnumeration;
import javax.naming.NamingException;
import javax.naming.directory.Attribute;
import javax.naming.directory.Attributes;
import javax.naming.directory.DirContext;
import javax.naming.directory.InitialDirContext;
import javax.naming.directory.SearchControls;
import javax.naming.directory.SearchResult;
import javax.security.auth.Subject;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.login.LoginException;

import org.apache.activemq.jaas.CertificateLoginModule;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

/**
 * This module builds on the ActiveMQ CertificateLoginModule to provide
 * directory based authorisation. The CertificateLoginModule's responsibility is
 * to obtain client-side certificates that are trusted by a server. This
 * CertificateLoginDirectoryRolesModule's responsibility is to take the subject
 * DN from the certificates and determine the roles (groups) that they are
 * members of. Note that there can be multiple client certificates presented.
 * Directory searches correlate one-to-one with the number of client
 * certificates.
 * 
 * <p>
 * Directory services are accessed via the JNDI directory API. If credentials
 * are provided (securityAuthentication != "none" i.e. "simple" or
 * <i>sasl_mech</i> etc.) then a bind occurs using the supplied
 * securityPrincipal and securityCredentials. Otherwise anonymous access is
 * performed. For more information: <a
 * href="http://java.sun.com/products/jndi/tutorial/ldap/security/ldap.html"
 * >LDAP Authentication</a>
 * 
 * <p>
 * What is searched for depends on searchBase as a starting point and
 * searchFilter. Search filters take one parameter denoted by the expression
 * {0}. The expression will be substituted with the subject DN of a client
 * certificate. Only the first level of the search base is searched. Search
 * filter syntax is described further at <a href=
 * "http://java.sun.com/products/jndi/tutorial/basics/directory/filter.html"
 * >Search Filters</a>
 * 
 * <p>
 * Here is a list of JAAS parameters that are required:
 * 
 * <table>
 * <tr>
 * <td>initialContextFactory</td>
 * <td>The fully distinguished JNDI context factory class name</td>
 * </tr>
 * <tr>
 * <td>providerURL</td>
 * <td>The URL required to connect to the directory</td>
 * </tr>
 * <tr>
 * <td>securityPrincipal</td>
 * <td>The name of the principal to bind to the directory. Can be omitted for
 * anonymous access.</td>
 * </tr>
 * <tr>
 * <td>securityCredentials</td>
 * <td>The password of the principal to bind to the directory. Can be omitted
 * for anonymous access.</td>
 * </tr>
 * <tr>
 * <td>securityAuthentication</td>
 * <td>"none", "simple" or <i>sasl_mech</i></td>
 * </tr>
 * <tr>
 * <td>searchBase</td>
 * <td>The part of the directory to perform one-level searches on.</td>
 * </tr>
 * <tr>
 * <td>searchFilter</td>
 * <td>The directory search specification. {0} will be substituted with each
 * certificate's subject DN.</td>
 * </tr>
 * <tr>
 * <td>searchReturnAttributeAsGroupName</td>
 * <td>The name of the attribute returned from the search that is used for the
 * name of the group.</td>
 * </tr>
 * </table>
 * 
 * Here is a sample JAAS login configuration that performs an anonymous search
 * of each group below &quot;ou=activemq,ou=groups,ou=system&quot. The
 * entry's "member" attribute equals the subject dn of each client certificate
 * presented to this module. The "cn" attribute is the part of the entry that is
 * used to provide the name of the group.
 * 
 * <pre>
 * activemq-domain {
 *   com.classactionpl.jaas.activemq.CertificateLoginDirectoryRolesModule required
 *     initialContextFactory=com.sun.jndi.ldap.LdapCtxFactory
 *     providerURL=&quot;ldap://localhost:10389&quot;
 *     securityAuthentication=none
 *     searchBase=&quot;ou=activemq,ou=groups,ou=system&quot;
 *     searchFilter=&quot;(member={0})&quot;
 *     searchReturnAttributeAsGroupName=&quot;cn&quot;
 * ;};
 * </pre>
 * 
 * @author Christopher Hunt
 * 
 */
public class CertificateLoginDirectoryRolesModule extends
		CertificateLoginModule {

	Log logger = LogFactory.getLog(CertificateLoginDirectoryRolesModule.class);

	Map<String, Set<String>> dnGroups = new HashMap<String, Set<String>>();

	Hashtable<Object, Object> env;

	String searchBase;
	String searchFilter;
	String searchReturnAttributeAsGroupName;

	@SuppressWarnings("unchecked")
	@Override
	public void initialize(Subject subject, CallbackHandler callbackHandler,
			Map sharedState, Map options) {
		// Build our jndi environment options from the options passed in
		String initialContextFactory = (String) options
				.get("initialContextFactory");
		String providerURL = (String) options.get("providerURL");
		String securityPrincipal = (String) options.get("securityPrincipal");
		String securityCredentials = (String) options
				.get("securityCredentials");
		String securityAuthentication = (String) options
				.get("securityAuthentication");
		searchBase = (String) options.get("searchBase");
		searchFilter = (String) options.get("searchFilter");
		searchReturnAttributeAsGroupName = (String) options
				.get("searchReturnAttributeAsGroupName");

		env = new Hashtable<Object, Object>();
		env.put(Context.INITIAL_CONTEXT_FACTORY, initialContextFactory);
		env.put(Context.PROVIDER_URL, providerURL);
		if (!securityAuthentication.equals("none")) {
			env.put(Context.SECURITY_PRINCIPAL, securityPrincipal);
			env.put(Context.SECURITY_CREDENTIALS, securityCredentials);
		}
		env.put(Context.SECURITY_AUTHENTICATION, securityAuthentication);

		// Now back to regular option processing
		super.initialize(subject, callbackHandler, sharedState, options);
	}

	@Override
	protected Set<String> getUserGroups(String dn) throws LoginException {
		return dnGroups.get(dn);
	}

	@Override
	protected String getUserNameForCertificates(X509Certificate[] certs)
			throws LoginException {
		dnGroups.clear();

		if (certs != null) {
			try {
				DirContext ctx = new InitialDirContext(env);

				try {
					for (X509Certificate cert : certs) {
						if (logger.isDebugEnabled()) {
							logger.debug("Client cert received: "
									+ cert.toString());
						}

						// Retrieve the groups that our subject dn is a member
						// of and keep track of them and their dn
						String dn = cert.getSubjectDN().getName();

						Set<String> groups = new HashSet<String>();
						dnGroups.put(dn, groups);

						SearchControls ctls = new SearchControls();
						ctls
								.setReturningAttributes(new String[] { searchReturnAttributeAsGroupName });
						NamingEnumeration<SearchResult> results = ctx.search(
								searchBase, searchFilter, new Object[] { dn },
								ctls);

						while (results.hasMore()) {
							SearchResult result = results.next();

							Attributes attrs = result.getAttributes();
							assert attrs != null;

							Attribute groupNameAttr = attrs.get("cn");
							assert groupNameAttr != null;

							String groupName = (String) groupNameAttr.get();
							groups.add(groupName);

							if (logger.isDebugEnabled()) {
								logger.debug("Found group: " + groupName);
							}
						}
					}
				} finally {
					ctx.close();
				}

			} catch (NamingException e) {
				logger.error(e.toString());
			}
		} else {
			logger.warn("No client certificates received");
		}
		return getDistinguishedName(certs);
	}
}
